from typing import Optional, Tuple, Union
import torch
from torch import Tensor

from torch.nn import Parameter
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.nn.dense.linear import Linear
from torch_geometric.nn.inits import zeros
from torch_geometric.typing import (
    Adj,
    OptPairTensor,
    OptTensor,
    SparseTensor,
)
from torch_geometric.utils import add_remaining_self_loops
from torch_geometric.utils import (
    scatter,
    spmm,
)
from torch_geometric.utils.num_nodes import maybe_num_nodes


def add_remaining_self_loops(edge_index, edge_attr=None, fill_value=None, num_nodes=None):
    print("==================增加自循环边=======================")
    print("原始 edge_index")
    print(edge_index)
    print("原始 edge_attr")
    print(edge_attr)
    N = num_nodes
    # 为源和目的不一样的边设定mask
    mask = edge_index[0] != edge_index[1]
    loop_index = torch.arange(0, N, dtype=torch.long, device=edge_index.device)
    loop_index = loop_index.unsqueeze(0).repeat(2, 1)
    # 建立一个新的二维矩阵，用于处理源和目的一致的情况
    if edge_attr is not None:
        if isinstance(fill_value, (int, float)):
            loop_attr = edge_attr.new_full((N,) + edge_attr.size()[1:],
                                           fill_value)
        inv_mask = ~mask
        loop_attr[edge_index[0][inv_mask]] = edge_attr[inv_mask]
        edge_attr = torch.cat([edge_attr[mask], loop_attr], dim=0)
        print(edge_attr)
    # 在没有边属性的情况
    edge_index = torch.cat([edge_index[:, mask], loop_index], dim=1)
    print("处理后 edge_index")
    print(edge_index)
    print("处理后 edge_attr")
    print(edge_attr)
    print("=========================================")
    return edge_index, edge_attr


def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False,
             add_self_loops=True, flow="source_to_target", dtype=None):
    fill_value = 2. if improved else 1.

    if add_self_loops:
        edge_index, edge_weight = add_remaining_self_loops(
            edge_index, edge_weight, fill_value, num_nodes)

    if edge_weight is None:
        # 边缘长度个1
        edge_weight = torch.ones((edge_index.size(1),), dtype=dtype,
                                 device=edge_index.device)
    print("边特征信息")
    print(edge_weight)
    print("=========================================")
    row, col = edge_index[0], edge_index[1]
    idx = col if flow == 'source_to_target' else row
    deg = scatter(edge_weight, idx, dim=0, dim_size=num_nodes, reduce='sum')
    print("度矩阵计算结果", deg)
    deg_inv_sqrt = deg.pow_(-0.5)
    deg_inv_sqrt.masked_fill_(deg_inv_sqrt == float('inf'), 0)
    print("度矩阵变换之后的内容", deg_inv_sqrt)
    # 左乘进行行变换，右乘进行列变换
    edge_weight = deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col]
    return edge_index, edge_weight


class GCNConv(MessagePassing):
    _cached_edge_index: Optional[OptPairTensor]
    _cached_adj_t: Optional[SparseTensor]

    def __init__(
            self,
            in_channels,
            out_channels,
            improved=False,
            cached=False,
            add_self_loops=True,
            normalize=True,
            bias=True,
            **kwargs,
    ):
        kwargs.setdefault('aggr', 'add')
        super().__init__(**kwargs)

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.improved = improved
        self.cached = cached
        self.add_self_loops = add_self_loops
        self.normalize = normalize

        self._cached_edge_index = None
        self._cached_adj_t = None

        self.lin = Linear(in_channels, out_channels, bias=False,
                          weight_initializer='glorot')

        if bias:
            self.bias = Parameter(torch.empty(out_channels))
        else:
            self.register_parameter('bias', None)

        self.reset_parameters()

    def reset_parameters(self):
        super().reset_parameters()
        self.lin.reset_parameters()
        zeros(self.bias)
        self._cached_edge_index = None
        self._cached_adj_t = None

    def forward(self, x: Tensor, edge_index: Adj,
                edge_weight: OptTensor = None) -> Tensor:

        if self.normalize:
            cache = self._cached_edge_index
            if cache is None:
                print("cache is None")
                print("edge_index", edge_index)
                print("edge_weight", edge_weight)
                print("x.size(self.node_dim) 节点数", x.size(self.node_dim))
                print("self.improved,", self.improved)
                print("self.add_self_loops, ", self.add_self_loops)
                print("self.flow", self.flow)
                print("x.dtype", x.dtype)
                # 做 GCN Normlize
                edge_index, edge_weight = gcn_norm(  # yapf: disable
                    edge_index, edge_weight, x.size(self.node_dim),
                    self.improved, self.add_self_loops, self.flow, x.dtype)
                if self.cached:
                    self._cached_edge_index = (edge_index, edge_weight)
            else:
                print(cache,"cache 333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333")
                edge_index, edge_weight = cache[0], cache[1]
        print("edge_index, edge_weight =====================================",edge_index, edge_weight)
        x = self.lin(x)
        # propagate_type: (x: Tensor, edge_weight: OptTensor)
        out = self.propagate(edge_index, x=x, edge_weight=edge_weight,
                             size=None)
        if self.bias is not None:
            out = out + self.bias
        return out

    def message(self, x_j: Tensor, edge_weight: OptTensor) -> Tensor:
        #对 x_j 进行更新 （发起端）
        return x_j if edge_weight is None else edge_weight.view(-1, 1) * x_j

    def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor:
        print("####################################################################################################################################################################################")
        print("######################### message_and_aggregate 处理前1")
        print(adj_t)
        print(x)
        print("######################### message_and_aggregate 处理前1")
        return spmm(adj_t, x, reduce=self.aggr)

    # def update(self,x_i):
    #     return x_i

edge_index = torch.tensor([[1, 2, 1, 2, 3, 3, 4], [2, 1, 0, 0, 0, 4, 3]], dtype=torch.long)  # 2 x E
x = torch.tensor([[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]], dtype=torch.float)  # N x emb(in)
edge_weight = torch.FloatTensor([[1], [2], [3], [4], [5], [6], [7]])


conv1 = GCNConv(2, 2,cached=False)
# conv1 = GCNConv(2, 2,cached=True)

out1 = conv1(x, edge_index, edge_weight)
out2 = conv1(out1, edge_index, edge_weight)
out3 = conv1(out2, edge_index, edge_weight)
print(out1)
print(out2)
print(out3)

